import tensorflow as tf
from tensorflow.python.util.nest import flatten
from tensorflow.keras import layers

x = tf.constant(
    [[
        [[1], [1], [0], [1], [0]],
        [[3], [-3], [4], [2], [1]],
        [[2], [0], [1], [3], [1]],
        [[4], [2], [-1], [0], [1]],
        [[6], [3], [-4], [1], [0]],
    ]], tf.float32
)
k = tf.constant(
    [[
        [[0]], [[1]], [[-1]]],
        [[[1]], [[0]], [[0]]],
        [[[1]], [[-1]], [[1]]
         ]], tf.float32
)
conv = tf.nn.conv2d(x, k, (1, 1, 1, 1), 'SAME')
print("conv:", conv)

avg_pool = tf.nn.avg_pool2d(conv, (2, 2), (1, 1, 1, 1), 'SAME')
print("avg_pool:", avg_pool)

max_pool = tf.nn.max_pool2d(conv, (2, 2), (1, 1, 1, 1), 'SAME')
print("max_pool:", max_pool)

flatten = flatten(max_pool)
print("flatten shape:", flatten)

# conv = [[
#     [[-6.][11.][-4.][3.][2.]]
#     [[-2.][7.][-2.][4.][4.]]
#     [[4.][-4.][5.][2.][3.]]
#     [[-1.][2.][8.][-4.][2.]]
#     [[2.][9.][2.][-5.][2.]]
# ]]
#
# avg_pool = [[
#     [[2.5][3.][0.25][3.25][3.]]
#     [[1.25][1.5][2.25][3.25][3.5]]
#     [[0.25][2.75][2.75][0.75][2.5]]
#     [[3.][5.25][0.25][-1.25][2.]]
#     [[5.5][5.5][-1.5][-1.5][2.]]
# ]]
#
# max_pool = [[
#     [[11.][11.][4.][4.][4.]]
#     [[7.][7.][5.][4.][4.]]
#     [[4.][8.][8.][3.][3.]]
#     [[9.][9.][8.][2.][2.]]
#     [[9.][9.][2.][2.][2.]]
# ]]
